// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/expression/field_reference.h>

#include <pollux/expression/peeled_encoding.h>

namespace kumo::pollux::exec {
    void FieldReference::computeDistinctFields() {
        SpecialForm::computeDistinctFields();
        if (inputs_.empty()) {
            mergeFields(
                distinctFields_,
                multiplyReferencedFields_,
                {this->as<FieldReference>()});
        }
    }

    void FieldReference::apply(
        const SelectivityVector &rows,
        EvalCtx &context,
        VectorPtr &result) {
        const RowVector *row;
        DecodedVector decoded;
        VectorPtr input;
        std::shared_ptr<PeeledEncoding> peeledEncoding;
        bool useDecode = false;
        LocalSelectivityVector nonNullRowsHolder(*context.execCtx());
        const SelectivityVector *nonNullRows = &rows;
        if (inputs_.empty()) {
            row = context.row();
        } else {
            inputs_[0]->eval(rows, context, input);

            decoded.decode(*input, rows);
            if (decoded.may_have_nulls()) {
                nonNullRowsHolder.get(rows);
                nonNullRowsHolder->deselectNulls(
                    decoded.nulls(&rows), rows.begin(), rows.end());
                nonNullRows = nonNullRowsHolder.get();
                if (!nonNullRows->hasSelections()) {
                    add_nulls(rows, decoded.nulls(&rows), context, result);
                    return;
                }
            }
            useDecode = !decoded.isIdentityMapping();
            if (useDecode) {
                std::vector<VectorPtr> peeledVectors;
                LocalDecodedVector localDecoded{context};
                peeledEncoding = PeeledEncoding::peel(
                    {input}, *nonNullRows, localDecoded, true, peeledVectors);
                POLLUX_CHECK_NOT_NULL(peeledEncoding);
                if (peeledVectors[0]->is_lazy()) {
                    peeledVectors[0] =
                            peeledVectors[0]->as<LazyVector>()->loaded_vector_shared();
                }
                POLLUX_CHECK(peeledVectors[0]->encoding() == VectorEncoding::Simple::ROW);
                row = peeledVectors[0]->as<const RowVector>();
            } else {
                POLLUX_CHECK(input->encoding() == VectorEncoding::Simple::ROW);
                row = input->as<const RowVector>();
            }
        }
        if (index_ == -1) {
            auto rowType = dynamic_cast<const RowType *>(row->type().get());
            POLLUX_CHECK(rowType);
            index_ = rowType->getChildIdx(field_);
        }
        VectorPtr child =
                inputs_.empty() ? context.getField(index_) : row->childAt(index_);
        if (child->encoding() == VectorEncoding::Simple::LAZY) {
            child = BaseVector::loaded_vector_shared(child);
        }
        if (result.get()) {
            if (useDecode) {
                child = peeledEncoding->wrap(type_, context.pool(), child, *nonNullRows);
            }
            result->copy(child.get(), *nonNullRows, nullptr);
        } else {
            // The caller relies on vectors having a meaningful size. If we
            // have a constant that is not wrapped in anything we set its size
            // to correspond to rows.end().
            if (!useDecode && child->is_constant_encoding()) {
                child = BaseVector::wrap_in_constant(nonNullRows->end(), 0, child);
            }
            result = useDecode
                         ? std::move(peeledEncoding->wrap(
                             type_, context.pool(), child, *nonNullRows))
                         : std::move(child);
        }
        child.reset();

        // Check for nulls in the input struct. Propagate these nulls to 'result'.
        if (!inputs_.empty() && decoded.may_have_nulls()) {
            add_nulls(rows, decoded.nulls(&rows), context, result);
        }
    }

    void FieldReference::evalSpecialForm(
        const SelectivityVector &rows,
        EvalCtx &context,
        VectorPtr &result) {
        VectorPtr localResult;
        apply(rows, context, localResult);
        context.moveOrCopyResult(localResult, rows, result);
    }

    void FieldReference::evalSpecialFormSimplified(
        const SelectivityVector &rows,
        EvalCtx &context,
        VectorPtr &result) {
        ExceptionContextSetter exceptionContext(
            {
                [](PolluxException::Type /*exceptionType*/, auto *expr) {
                    return static_cast<Expr *>(expr)->toString();
                },
                this
            });
        VectorPtr input;
        const RowVector *row;
        if (inputs_.empty()) {
            row = context.row();
        } else {
            POLLUX_CHECK_EQ(inputs_.size(), 1);
            inputs_[0]->evalSimplified(rows, context, input);
            BaseVector::flattenVector(input);
            row = input->as<RowVector>();
            POLLUX_CHECK(row);
        }
        auto index = row->type()->as_row().getChildIdx(field_);
        if (index_ == -1) {
            index_ = index;
        } else {
            POLLUX_CHECK_EQ(index_, index);
        }

        LocalSelectivityVector nonNullRowsHolder(*context.execCtx());
        const SelectivityVector *nonNullRows = &rows;
        if (row->may_have_nulls()) {
            nonNullRowsHolder.get(rows);
            nonNullRowsHolder->deselectNulls(row->raw_nulls(), rows.begin(), rows.end());
            nonNullRows = nonNullRowsHolder.get();
            if (!nonNullRows->hasSelections()) {
                add_nulls(rows, row->raw_nulls(), context, result);
                return;
            }
        }

        auto &child = row->childAt(index_);
        context.ensure_writable(rows, type_, result);
        result->copy(child.get(), *nonNullRows, nullptr);
        if (row->may_have_nulls()) {
            add_nulls(rows, row->raw_nulls(), context, result);
        }
    }

    std::string FieldReference::toString(bool recursive) const {
        std::stringstream out;
        if (!inputs_.empty() && recursive) {
            appendInputs(out);
            out << ".";
        }
        out << name();
        return out.str();
    }

    std::string FieldReference::toSql(
        std::vector<VectorPtr> *complexConstants) const {
        std::stringstream out;
        if (!inputs_.empty()) {
            appendInputsSql(out, complexConstants);
            out << ".";
        }
        out << "\"" << name() << "\"";
        return out.str();
    }
} // namespace kumo::pollux::exec
